"""
该函数设计出来用于判断两个正方体在经过RT变换之后，每个顶点之间的距离，并返回其最大值
新增 2024-04-26 17:23:03
在计算MPD时可以选择计算顶点均值
"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.spatial.transform import Rotation as R
import json
import pandas as pd
from DRR.ct_projector import add_noise, euler_angles2rot_matrix


def cube_max_distance(rt1, rt2, render=False, mode='max'):
    # 建立一个180*180*100的矩阵，并将每个顶点的位置信息储存
    cube = np.array([[-90, 90, -50, 1],
                     [-90, -90, -50, 1],
                     [90, 90, -50, 1],
                     [90, -90, -50, 1],
                     [-90, 90, 50, 1],
                     [-90, -90, 50, 1],
                     [90, 90, 50, 1],
                     [90, -90, 50, 1]])
    new_pos1 = np.dot(rt1, cube.T).T
    new_pos2 = np.dot(rt2, cube.T).T
    # 将三维点位显示出来
    if render:
        ax = plt.figure().add_subplot(111, projection='3d')
        ax.scatter(cube[:, 0], cube[:, 1], cube[:, 2], c='g', marker='o', label='origin')
        ax.scatter(new_pos1[:, 0], new_pos1[:, 1], new_pos1[:, 2], c='b', marker='o', label='pos1')
        ax.scatter(new_pos2[:, 0], new_pos2[:, 1], new_pos2[:, 2], c='r', marker='o', label='pos1')
        ax.legend()
        plt.show()
    # 计算两个新位置之间顶点的距离
    distance = np.linalg.norm((new_pos1 - new_pos2), axis=1)
    if mode == 'average':
        return np.mean(distance)
    else:
        return max(distance)


def get_default_rt(rot_cen=np.array([89.82421875, 89.82421875, 53]), d_s2c=400,
                   theta=None, beta=None, alpha=None):
    """
    compute the Rt matrix from x-ray source coordinate to CT coordinate

    :param rot_cen: rotation center of C-Arm (mm) in CT coordinate
    :param d_s2c: distance between x-ray source and rot_cen (mm)
    :param theta: rotation angle along z axis of the CT data
    :param beta: angle between z axis of CT and z axis of x-ray source, i.e. rotation angle along y-axis of x-ray source
    :param alpha: rotation angle of image plane along z axis of the C-Arm
    :return: transform matrix from x-ray source coordinate to CT coordinate
    """
    rot_z = euler_angles2rot_matrix(0, 0, alpha * np.pi / 180)
    # compute rotation matrix from Euler angles (theta_x, theta_y, theta_z) with the default rotation order
    rot_matrix = euler_angles2rot_matrix(0, np.pi * beta / 180, theta * np.pi / 180) @ rot_z
    # translation vector i.e. the 3D coordinates of the x-ray source in the CT coordinates system
    t_vec = np.array(rot_matrix @ np.mat([0, 0, -d_s2c]).T).squeeze() + np.array(rot_cen)
    return np.column_stack((rot_matrix, t_vec)).astype(np.float32)


def get_rt(alpha_noise, beta_noise, theta_noise, tx_noise, ty_noise, tz_noise,
           mode='标准正位', rot_cen=np.array([89.82421875, 89.82421875, 53]), d_s2c=400):
    Rt = None
    # 先得到默认的rt
    if mode == "标准正位":
        # 标准正位：（0±5， 270±10， 90±5）, t(0±25， 0±50， 0±25)
        Rt = get_default_rt(alpha=0, beta=270, theta=90, rot_cen=rot_cen, d_s2c=d_s2c)
    Rt = np.vstack((Rt, [0, 0, 0, 1]))
    Rt = (add_noise(np.array([alpha_noise, beta_noise, theta_noise]).flatten(),
                    np.array([tx_noise, ty_noise, tz_noise]).flatten(),
                    rot_cen=rot_cen) @ Rt)
    return Rt


if __name__ == "__main__":
    # RT1 = get_rt(90, 90, 90, 10, 0, 0)
    # RT2 = get_rt(0, 0, 0, 5, 0, 0)
    # max_dis = cube_max_distance(RT1, RT2, True)
    # print(max_dis)
    # 读取json文件
    # im_dict = {}
    # with open('train_label_info.json', 'r') as test_json:
    #     test_info = json.load(test_json)
    # for key in test_info:
    #     print(test_info[key])
    #     # 字典文件保存文件标签
    #     im_dict[key] = {
    #         "name": test_info[key]['name'],
    #         "angle_x": test_info[key]['angle_x'],
    #         "angle_y": test_info[key]['angel_y'],
    #         "angle_z": test_info[key]['angel_z'],
    #         "tx": test_info[key]['tx'],
    #         "ty": test_info[key]['ty'],
    #         "tz": test_info[key]['tz'],
    #     }
    # with open("D:/大四/毕设（实际工作）/六维图像配准最终版/train_label_info2.json", "w") as f:
    #     json.dump(im_dict, f)
    # 任意的多组列表
    a = [1, 2, 3]
    b = [4, 5, 6]

    # 字典中的key值即为csv中列名
    dataframe = pd.DataFrame({'epoch': [], 'loss_per_20': [],
                              'test_loss_per_20': [], 'distance_per_20': [],
                              'test_distance_per_20': [], 'time': []})

    # 将DataFrame存储为csv,index表示是否显示行名，default=True
    dataframe.to_csv("6D_train_history/train_log.csv", index=False, sep=',')
